from PIL import Image
import matplotlib.pyplot as plt
import torchvision
import cv2
import os
import re
import numpy as np
import random

pic_size = 384

class RandomCropResizeLongestSide:
    def __init__(self, max_length=pic_size, scale=(0.7, 1.2), ratio=(0.8, 1.2)):
        self.max_length = max_length
        self.scale = scale
        self.ratio = ratio

    def __call__(self, img):
        # Randomly crop the image
        i, j, h, w = torchvision.transforms.RandomResizedCrop.get_params(img, self.scale, self.ratio)
        img = torchvision.transforms.functional.crop(img, i, j, h, w)

        # Calculate the new size while maintaining the aspect ratio
        if h > w:
            new_h = self.max_length
            new_w = int(new_h * w / h)
        else:
            new_w = self.max_length
            new_h = int(new_w * h / w)

        # Resize the image to the new size
        img = torchvision.transforms.functional.resize(img, (new_h, new_w))

        return img

def show_images(imgs, nrows, ncols, scale=2):
    _, axes = plt.subplots(nrows, ncols, figsize=(nrows * scale, ncols * scale))
    for i in range(nrows):
        for j in range(ncols):
            axes[i, j].imshow(imgs[i * ncols + j])
            axes[i, j].axes.get_xaxis().set_visible(True)
            axes[i, j].axes.get_yaxis().set_visible(False)
    plt.show()
    return axes

def apply(img, auc, nrows=2, ncols=4, scale=1.5):
    Y = [auc(img) for _ in range(nrows * ncols)]
    # for i in Y:
        # i.show()
    show_images(Y, nrows, ncols, scale)


augs = torchvision.transforms.Compose(
    [
        # torchvision.transforms.RandomErasing( scale=(0.02, 0.33), ratio=(0.3, 3.3), value='random',inplace=True),
        torchvision.transforms.RandomHorizontalFlip(p=0.5),
        RandomCropResizeLongestSide(max_length=pic_size, scale=(0.8, 1.2), ratio=(0.8, 1.2)),
        # torchvision.transforms.RandomResizedCrop(pic_size, scale=(0.7, 1.2), ratio=(0.8, 1.2)),
        torchvision.transforms.RandomRotation(8),
        torchvision.transforms.ColorJitter(brightness=0, hue=0.4, contrast=0, saturation=0.4)
    ]
)


def traverse_dir(path):
    file_dir=[]
    for root, dirs, files in os.walk(path):
        for x in files:
            file_dir.append(path+'\\'+x)
            # print(path+'\\'+x)
    return file_dir

if __name__ == '__main__':
    path = r"G:\python\data\sdfcar\train"
    # img = Image.open(r'E:\mynet\resnet18\data\train\2\02314.jpg')
    # # img.show()
    # s = augss(img)
    # s.show()
    for i in range(1,196):
        index = 0
        paths = path + "\\"+str(i)
        for pic in traverse_dir(paths):
            # print(pic)
            have = pic.find('a.')
            if(have == -1):
                img = Image.open(pic)
                for i in range(1):
                    # print(pic)
                    s= augs(img)
                    s.save(paths + "\\" + str(index) + "a.jpg")
                    index =index +1
                # s.show()

    # img.show()
    # apply(img, torchvision.transforms.RandomHorizontalFlip(p=0.5))
    # apply(img, torchvision.transforms.RandomResizedCrop(300, scale=(0.5, 1), ratio=(0.6, 1.6)))